from common_import import *
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import spearmanr


def spearman_analysis():
    data = []
    for i in range(1, 7):  # 假设有6个产品
        product_data = tool.get_np(f"type_monthly_total/{i}.csv")["monthly_total"]
        data.append(product_data)

    data = np.array(data).T

    # 计算斯皮尔曼相关系数矩阵
    spearman_corr_matrix, _ = spearmanr(data, axis=0)

    print("斯皮尔曼相关系数矩阵：")
    print(spearman_corr_matrix)

    # 类别名称
    labels = mapping.type_map1.values()

    # 创建热力图
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        spearman_corr_matrix,
        annot=True,
        cmap="coolwarm",
        xticklabels=labels,
        yticklabels=labels,
        vmin=-1,
        vmax=1,
        cbar_kws={"label": ""},
    )
    plt.yticks(rotation=0)
    # 添加标题
    plt.title("斯皮尔曼相关系数矩阵")

    # 显示图表
    mydraw.show_or_print("spearman_corr.png")


def spearman_corr_product():
    data = []
    pro_list = [5, 6, 7, 10, 15, 16, 19]
    for i in pro_list:  # 假设有7个产品
        product_data = tool.get_np(f"product_monthly_total/{i}.csv")["monthly_total"]
        data.append(product_data)

    data = np.array(data).T

    # 计算斯皮尔曼相关系数矩阵
    spearman_corr_matrix, _ = spearmanr(data, axis=0)

    # 打印斯皮尔曼相关系数矩阵
    print(spearman_corr_matrix)

    # 类别名称
    labels = [mapping.get_name(i) for i in pro_list]

    # 创建热力图
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        spearman_corr_matrix,
        annot=True,
        cmap="coolwarm",
        xticklabels=labels,
        yticklabels=labels,
        vmin=-1,
        vmax=1,
        cbar_kws={"label": ""},
    )
    plt.yticks(rotation=0)

    # 添加标题
    plt.title("斯皮尔曼相关系数矩阵(产品)")

    # 显示或保存图表
    mydraw.show_or_print("spearman_corr_product.png")


if __name__ == "__main__":
    spearman_corr_product()
